#include "elemental.hpp"
#include <omp.h>

using namespace std;
using namespace elem;


void Usage()
{
    cout << "GEneral Matrix Matrix multiplication.\n\n"
         << "  Gemm <r> <c> <orient. of A?> <orient. of B?> <m> <n> <k> <nb> " 
            "<print?>\n\n"
         << "  r: number of process rows\n"
         << "  c: number of process cols\n"
         << "  m: height of C\n" 
         << "  n: width  of C\n"
         << "  k: inner dimension of AB\n"
         << "  nb: algorithmic blocksize\n"
         << "  print?: [0/1]\n"
         << "  b: buffer size\n"
         << "  verify?: [0/1] (REQUIRES VERY LARGE BUFFERS !!) \n" 
         << "  runs: number of runs to loop\n" << endl; 
}

template<typename T> // represents a real or complex ring
void TestGemm
( bool verify, bool printMatrices, Orientation orientationOfA, Orientation orientationOfB,
  int m, int n, int k, T alpha, T beta, const Grid& g )
{
    double startTime, endTime, runTime, gFlops;
    DistMatrix<T> A(g), B(g), C(g);
    DistMatrix<T, STAR, STAR> A_STAR_STAR(g);
    DistMatrix<T, STAR, STAR> B_STAR_STAR(g);
    DistMatrix<T, STAR, STAR> C_STAR_STAR(g);
    if( orientationOfA == NORMAL )
        A.ResizeTo( m, k );
    else
        A.ResizeTo( k, m );
    if( orientationOfB == NORMAL )
        B.ResizeTo( k, n );
    else
        B.ResizeTo( n, k );
    C.ResizeTo( m, n );
    A_STAR_STAR.ResizeTo( m, k );
    B_STAR_STAR.ResizeTo( k, n );
    C_STAR_STAR.ResizeTo( m, n );
    if( g.Rank() == 0 )
        cout << endl << "Stationary C Algorithm:" << endl;
    
    if( g.Rank() == 0 )
        cout << endl << "Make A:" << endl;
    MakeUniform( A );
    if( g.Rank() == 0 )
        cout << endl << "Make B:" << endl;
    MakeUniform( B );
    if( g.Rank() == 0 )
        cout << endl << "Make C:" << endl;
    MakeUniform( C );
    
    if ( verify ) {
    
        A_STAR_STAR = A;
        B_STAR_STAR = B;
        C_STAR_STAR = C;
    
    }
    if( printMatrices )
    {
        A.Print("A");
        B.Print("B");
        C.Print("C");
    }
    if( g.Rank() == 0 )
    {
        cout << "  Starting Gemm...";
        cout.flush();
    }
    
#pragma omp barrier

    startTime = omp_get_wtime();
#ifndef RELEASE
    PushCallStack("============= GEMM =============");
#endif
    internal::GemmC
    ( orientationOfA, orientationOfB, alpha, A, B, beta, C );
#ifndef RELEASE
    PushCallStack("============= GEMM =============");
#endif

#pragma omp barrier
    
    endTime = omp_get_wtime();
    runTime = endTime - startTime;
    gFlops = internal::GemmGFlops<T>(m,n,k,runTime);
    if( g.Rank() == 0 )
    {
        cout << "DONE. " << endl
             << "  Time = " << runTime << " seconds. GFlops = " 
             << gFlops << endl;
    }
    if ( verify ) {
    if( g.Rank() == 0 ) {
        startTime = omp_get_wtime();
        internal::LocalGemm
            ( NORMAL, NORMAL, (T) alpha, A_STAR_STAR, B_STAR_STAR, (T) beta, C_STAR_STAR );
        endTime = omp_get_wtime();
        
        runTime = endTime - startTime;
        gFlops = internal::GemmGFlops<T>(m,n,k,runTime);
        
        
            cout << "DONE. " << endl
                << "  Time = " << runTime << " seconds. GFlops = " 
                << gFlops << endl;
        }
    }
    
    
    if( printMatrices )
    {
        ostringstream msg;
        msg << "C ( *, *):= " << alpha << " A ( *, *) B ( *, *) + " << beta << " C";
        C.Print( msg.str() );
        
        if ( verify ) {
        
        msg << "( *, *)";
        C_STAR_STAR.Print( msg.str() );
        
        }
    }
    
    if ( verify ) {
        A_STAR_STAR = C;
        MakeIdentity( B_STAR_STAR);
        A_STAR_STAR.Print( "C 'Distributed'");
        B_STAR_STAR.Print( "I");
        C_STAR_STAR.Print( "C ( *, *) 'Single Core'");
        
        internal::LocalGemm
            ( NORMAL, NORMAL, (T) 1, A_STAR_STAR, B_STAR_STAR, (T) -1, C_STAR_STAR );
        
        
        
        if( printMatrices )
        {
            ostringstream msg;
            msg << "ZEROS C ( *, *) = 1 C * I - 1 C ( *, *)";
            C_STAR_STAR.Print( msg.str() );
        }
        
        if (omp_get_thread_num() == 0) {
            std::cout << "SUM:" << C_STAR_STAR.SumMatrix() << std::endl;
        }
    }
}

void TestGrid( mpi::Comm comm, const char* msg) {
    
#pragma omp barrier
    int w_rank = omp_get_thread_num();
    if (w_rank == 0) std::cout << msg << std::endl;
    
#pragma omp barrier
        
    int mpi_size = omp_get_num_threads();
    std::vector<int> comm_ = GetComm(comm.threadId, comm.type);
    int comm_size = comm_.size();
    int comm_rank = -1;
    
    std::stringstream aux;
    for (int i = 0; i < comm_size; i++) {
        if ( comm_[i] == w_rank) comm_rank = i;
        aux << comm_[i] << " ";
    }
    aux << "\n";
    
    
    for(int i = 0; i < mpi_size; i++) {
        if (i == w_rank) {
            
            std::stringstream ss;
            ss << "comm_tid:" << comm_rank 
                << "\ntid:" << w_rank
                << "\nsize:" << comm_size
                << "\n";
            
            std::cout << ss.str();
            //std::cout << aux.str();
        }
        
        
#pragma omp barrier
    }
#pragma omp barrier
}

int main( int argc, char** argv ) {
    
    if( argc < 11 )
    {
        Usage();
        return 0;
    }
    
    int argNum = 0;
    const int r = atoi(argv[++argNum]);
    const int c = atoi(argv[++argNum]);
    const Orientation orientationOfA = CharToOrientation('N');
    const Orientation orientationOfB = CharToOrientation('N');
    const int m = atoi(argv[++argNum]);
    const int n = atoi(argv[++argNum]);
    const int k = atoi(argv[++argNum]);
    const int nb = atoi(argv[++argNum]);
    const bool printMatrices = atoi(argv[++argNum]);
    const unsigned long b = atol(argv[++argNum]);
    const bool verify = atoi(argv[++argNum]);
    const int runs = atoi(argv[++argNum]);
    
    //const unsigned long buffer = 100000000;
    
    Initialize( r, c, b);
    
    
    omp_set_num_threads(r * c);
    
    
#pragma omp parallel
    {
        const int rank = omp_get_thread_num();
#ifndef RELEASE
        if( rank == 0 )
        {
            cout << "==========================================\n"
                 << " In debug mode! Performance will be poor! \n"
                 << "==========================================" << endl;
        }
#endif



#pragma omp barrier
        const Grid g( r, c);
#pragma omp barrier
        SetBlocksize( nb );
#pragma omp barrier
        

        
        /*
         * TEST GRID DATA
         *
         */
   
 //TestGrid( g.VCComm() , "VCComm");
 //TestGrid( g.VRComm() , "VRComm");
 //TestGrid( g.MCComm() , "MCComm");
 //TestGrid( g.MRComm() , "MRComm");

        /*
         * END TEST GRID DATA
         */

        
        if( rank == 0 )
        {
            cout << "Will test Gemm" << OrientationToChar(orientationOfA) 
                                     << OrientationToChar(orientationOfB) 
                                     << endl;
        }
        
        if( rank == 0 )
        {
            cout << "---------------------\n"
                 << "Testing with doubles:\n"
                 << "---------------------" << endl;
        }
        for (int i = 0; i < runs; i++)
            TestGemm<double> ( verify, printMatrices, orientationOfA, orientationOfB, m, n, k, (double)3, (double)2, g );
    }
    
    Finalize();
#ifndef RELEASE
    //DumpCallStack();
#endif
}
    